#!/usr/bin/env python3
"""
集成式多层记忆架构AI系统 - 主入口文件
支持HTTP和Stdio两种MCP连接方式
"""

import asyncio
import sys
import argparse
from pathlib import Path

from fastapi import FastAPI
from loguru import logger

from src.config import load_config
from src.mcp.server import MCPServer
from src.core.context_orchestrator import ContextOrchestrator
from src.database.manager import DatabaseManager


async def start_http_server(config: dict):
    """启动HTTP服务器"""
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware
    
    # 创建FastAPI应用
    app = FastAPI(
        title="MCP Server",
        description="集成式多层记忆架构AI系统",
        version="1.0.0"
    )
    
    # 添加CORS中间件
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    @app.get("/")
    async def root():
        """根路径"""
        return {"message": "MCP Server is running", "version": "1.0.0"}
    
    @app.get("/health")
    async def health_check():
        """健康检查"""
        return {"status": "healthy"}
    
    import uvicorn
    server_config = uvicorn.Config(
        app,
        host=config['mcp']['http']['host'],
        port=config['mcp']['http']['port'],
        log_level="info",
        access_log=True
    )
    
    server = uvicorn.Server(server_config)
    logger.info(f"启动HTTP服务器: {config['mcp']['http']['host']}:{config['mcp']['http']['port']}")
    await server.serve()


async def start_stdio_server(config: dict):
    """启动Stdio服务器"""
    from src.mcp.stdio_handler import StdioHandler
    
    # 创建一个简化的MCP服务器，先不依赖数据库
    # 这样可以立即响应MCP协议消息
    mcp_server = MCPServer(None, config)
    
    # 创建Stdio处理器
    stdio_handler = StdioHandler(mcp_server, config)
    
    logger.info("启动Stdio MCP服务器")
    
    # 在后台异步初始化数据库和上下文编排器
    async def initialize_backend():
        try:
            logger.info("开始后台初始化数据库和上下文编排器...")
            
            # 初始化数据库管理器
            db_manager = DatabaseManager(config)
            await db_manager.initialize()
            
            # 初始化上下文编排器
            orchestrator = ContextOrchestrator(config)
            await orchestrator.initialize()
            
            # 更新MCP服务器的编排器
            mcp_server.orchestrator = orchestrator
            
            logger.info("后台初始化完成，所有功能现已可用")
            
        except Exception as e:
            logger.error(f"后台初始化失败: {e}")
    
    # 启动后台初始化任务
    asyncio.create_task(initialize_backend())
    
    # 立即启动stdio处理器
    await stdio_handler.start()


async def start_websocket_server(config: dict):
    """启动WebSocket服务器（用于Stdio over WebSocket）"""
    from src.mcp.websocket_handler import WebSocketHandler
    
    # 初始化数据库管理器
    db_manager = DatabaseManager(config)
    await db_manager.initialize()
    
    # 初始化上下文编排器
    orchestrator = ContextOrchestrator(config)
    await orchestrator.initialize()
    
    # 创建MCP服务器
    mcp_server = MCPServer(orchestrator, config)
    
    # 创建WebSocket处理器
    ws_handler = WebSocketHandler(mcp_server, config)
    
    logger.info(f"启动WebSocket服务器: {config['mcp']['websocket']['host']}:{config['mcp']['websocket']['port']}")
    await ws_handler.start()


def setup_logging(config: dict, mode: str = 'http'):
    """配置日志"""
    log_config = config.get('logging', {})
    
    # 移除默认处理器
    logger.remove()
    
    # 在stdio模式下，必须确保所有日志都输出到stderr，避免污染stdout
    if mode == 'stdio':
        # stdio模式：只输出到stderr和文件，避免污染stdout
        logger.add(
            sys.stderr,
            level=log_config.get('level', 'ERROR'),  # stdio模式下减少日志级别
            format='{time} | {level} | {name}:{function}:{line} - {message}',
            colorize=False  # 避免ANSI转义序列
        )
    else:
        # 其他模式：正常的控制台处理器
        logger.add(
            sys.stderr,
            level=log_config.get('level', 'INFO'),
            format=log_config.get('format', '{time} | {level} | {message}'),
            colorize=True
        )
    
    # 添加文件处理器
    if log_config.get('file_path'):
        log_file = Path(log_config['file_path'])
        log_file.parent.mkdir(parents=True, exist_ok=True)
        
        logger.add(
            log_file,
            level=log_config.get('level', 'INFO'),
            format=log_config.get('format', '{time} | {level} | {message}'),
            rotation=log_config.get('rotation', '1 day'),
            retention=log_config.get('retention', '30 days'),
            compression=log_config.get('compression', 'gz')
        )


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='集成式多层记忆架构AI系统')
    parser.add_argument(
        '--mode', 
        choices=['http', 'stdio', 'websocket'], 
        default='http',
        help='服务器模式: http(默认), stdio, websocket'
    )
    parser.add_argument(
        '--config', 
        default='config',
        help='配置文件目录路径'
    )
    
    args = parser.parse_args()
    
    # 加载配置
    try:
        config = load_config(args.config)
        setup_logging(config, args.mode)
        logger.info(f"配置加载成功，启动模式: {args.mode}")
    except Exception as e:
        print(f"配置加载失败: {e}")
        sys.exit(1)
    
    # 根据模式启动相应服务器
    try:
        if args.mode == 'http':
            await start_http_server(config)
        elif args.mode == 'stdio':
            await start_stdio_server(config)
        elif args.mode == 'websocket':
            await start_websocket_server(config)
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在关闭服务器...")
    except Exception as e:
        logger.error(f"服务器启动失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())